RFC: Make it easier to call window functions via expression API (and add example) #6746

alamb · 2023-06-22T17:09:40Z

before polishing this PR up I wanted to get some feedback on the approach

Which issue does this PR close?

Related to #5781
Closes: #6747

Rationale for this change

@mustafasrepo suggested adding an example for DataFrame::window here: #6703 (comment)

When I tried to do so it turns out that creating an Expr::Window is quite challenging, so I took a step back and tried to clean up the API a little

What changes are included in this PR?

Change WindowExpr to be a builder style interface and thus making it easier to construct Expr::WindowFunction variants. This is modeled after the CaseBuilder
Add the lead, lag, etc. in expr_fns to follow the model of the other types of functions

Are these changes tested?

Yes

TOOD: I need to write tests for all the code in expr_fns which I will add to https://github.com/apache/arrow-datafusion/blob/main/datafusion/core/tests/dataframe/dataframe_functions.rs if people like this basic approach

Are there any user-facing changes?

yes, the APIs to create window functions are improved

alamb · 2023-06-22T17:10:21Z

datafusion/core/src/dataframe.rs

+    /// let df = ctx.read_csv("tests/data/example.csv", CsvReadOptions::new()).await?;
+    ///
+    /// // The following is the equivalent of "SELECT FIRST_VALUE(b) OVER(PARTITION BY a)"
+    /// let first_value = first_value(col("b"))


The rest of the PR is motivated by trying to write this example. It turned out to be quite a pain to create the right Expr.

alamb · 2023-06-22T17:11:29Z

datafusion/expr/src/expr.rs

-    ) -> Self {
+    /// Create a new Window expression with the specified argument an
+    /// empty `OVER` clause
+    pub fn new(fun: impl Into<window_function::WindowFunction>, args: Vec<Expr>) -> Self {


This is the major change -- so that WindowFunction acts like a Builder which can be used to set partition by, order by and window_frame with a more self documenting style

timsaucer · 2024-05-02T13:06:43Z

As a new user, I like this change very much. I found creating lead and lag functions to be difficult to use. Suppose I wanted to call lead but provide the other two optional arguments - shift offset and default value. Would the best way to do that with this proposal to create the WindowFunction directly?

let single_lead = lead(col("mycol"));
let two_lead_with_default = expr::WindowFunction::new(BuiltInWindowFunction::Lead, vec![col("mycol"), lit(2), lit(mydefault)]);

alamb · 2024-05-03T09:49:59Z

As a new user, I like this change very much. I found creating lead and lag functions to be difficult to use. Suppose I wanted to call lead but provide the other two optional arguments - shift offset and default value. Would the best way to do that with this proposal to create the WindowFunction directly?

I think it would look something like this

let single_lead = lead(col("mycol"))
  .
let two_lead_with_default = lead.call(col("mycol"))
  .offset(lit(2))
  .partition_by(lit(mydefault)]))
  .build()

timsaucer · 2024-05-03T12:00:19Z

I'm looking at the arguments we end up passing. I think cume_dist takes no arguments. Lead and lag take anywhere between 1 and 3. nth_value takes two. I think it might be easier to use with something like

pub fn lead(arg: Expr, offset: Option<i64>, default: Option<ScalarValue>) -> expr::WindowFunction {
    let offset = lit(ScalarValue::Int64(offset));
    let default_value = match default {
        Some(v) => lit(v),
        None => lit(ScalarValue::Null),
    };
    expr::WindowFunction::new(BuiltInWindowFunction::Lead, vec![arg, offset, default_value])
}

and for nth value it's easier

pub fn nth_value(arg: Expr, n: i64) -> expr::WindowFunction {
    expr::WindowFunction::new(BuiltInWindowFunction::NthValue, vec![arg, lit(n)])
}

alamb · 2024-05-06T16:39:53Z

I think it might be easier to use with something like

That makes sense to me @timsaucer -- thank you for the feedback.

I don't think I will have time to work on this PR / issue for a while, but I tried to give it some visibility in #10395 and hopefully we can find someone else who is interested in helping make it work

alamb · 2024-05-06T16:40:09Z

(or of course if you have the time it would be an amazing contribution 🙏 )

timsaucer · 2024-05-07T16:10:19Z

I'm willing to work on this, but I'd like to wrap up the examples I have going into datafusion-python first. I think those will have more impact. I only have an hour or so a day to work on this, so I can't promise how fast I'll get to it.

alamb · 2024-05-07T19:16:33Z

I'm willing to work on this, but I'd like to wrap up the examples I have going into datafusion-python first. I think those will have more impact. I only have an hour or so a day to work on this, so I can't promise how fast I'll get to it.

Thank you very much 🙏

I totally understand re time constraints. Maybe we'll find some other people to help too. One can hope

github-actions · 2024-07-07T01:53:43Z

Thank you for your contribution. Unfortunately, this pull request is stale because it has been open 60 days with no activity. Please remove the stale label or comment or this will be closed in 7 days.

alamb · 2024-07-07T11:30:07Z

We can revive this PR / its API when someone has time to work on it

In case anyone is following along, @jayzhan211 added a really nice trait for working with aggregate functions. Maybe we can do something similar for window functions eventually

datafusion/datafusion/expr/src/udaf.rs

Lines 614 to 653 in e693ed7

    
           /// Extensions for configuring [`Expr::AggregateFunction`] 
        
           /// 
        
           /// Adds methods to [`Expr`] that make it easy to set optional aggregate options 
        
           /// such as `ORDER BY`, `FILTER` and `DISTINCT` 
        
           /// 
        
           /// # Example 
        
           /// ```no_run 
        
           /// # use datafusion_common::Result; 
        
           /// # use datafusion_expr::{AggregateUDF, col, Expr, lit}; 
        
           /// # use sqlparser::ast::NullTreatment; 
        
           /// # fn count(arg: Expr) -> Expr { todo!{} } 
        
           /// # fn first_value(arg: Expr) -> Expr { todo!{} } 
        
           /// # fn main() -> Result<()> { 
        
           /// use datafusion_expr::AggregateExt; 
        
           /// 
        
           /// // Create COUNT(x FILTER y > 5) 
        
           /// let agg = count(col("x")) 
        
           ///    .filter(col("y").gt(lit(5))) 
        
           ///    .build()?; 
        
           ///  // Create FIRST_VALUE(x ORDER BY y IGNORE NULLS) 
        
           /// let sort_expr = col("y").sort(true, true); 
        
           /// let agg = first_value(col("x")) 
        
           ///   .order_by(vec![sort_expr]) 
        
           ///   .null_treatment(NullTreatment::IgnoreNulls) 
        
           ///   .build()?; 
        
           /// # Ok(()) 
        
           /// # } 
        
           /// ``` 
        
           pub trait AggregateExt { 
        
               /// Add `ORDER BY <order_by>` 
        
               /// 
        
               /// Note: `order_by` must be [`Expr::Sort`] 
        
               fn order_by(self, order_by: Vec<Expr>) -> AggregateBuilder; 
        
               /// Add `FILTER <filter>` 
        
               fn filter(self, filter: Expr) -> AggregateBuilder; 
        
               /// Add `DISTINCT` 
        
               fn distinct(self) -> AggregateBuilder; 
        
               /// Add `RESPECT NULLS` or `IGNORE NULLS` 
        
               fn null_treatment(self, null_treatment: NullTreatment) -> AggregateBuilder; 
        
           }

timsaucer · 2024-07-08T12:26:04Z

I should have time to work on this near the end of the week or early next week. I'm in nearing the end of a large PR for datafusion-python and can dive back on this topic in once that's up for review.

alamb · 2024-07-09T21:38:10Z

Thanks @timsaucer -- looking forward to it.

I think the most useful part of this PR are the doc examples

Make it easier to call window functions via expression API

18f06ae

github-actions bot added core Core DataFusion crate logical-expr Logical plan and expressions optimizer Optimizer rules sql SQL Planner labels Jun 22, 2023

alamb commented Jun 22, 2023

View reviewed changes

alamb mentioned this pull request Jun 22, 2023

Support User Defined Window Functions #6703

Merged

tustvold force-pushed the main branch from 9dd3c2b to 70d633a Compare June 27, 2023 09:44

tustvold changed the base branch from main to backup-main June 27, 2023 11:12

tustvold changed the base branch from backup-main to main June 27, 2023 11:13

alamb self-assigned this Nov 6, 2023

alamb mentioned this pull request Nov 15, 2023

Add Library Guide for User Defined Functions: Window/Aggregate #8171

Merged

alamb mentioned this pull request May 2, 2024

Make it easier to create WindowFunctions with the Expr API #6747

Closed

alamb mentioned this pull request May 6, 2024

DataFusion weekly project plan (Andrew Lamb) - May 6, 2024 #10395

Closed

7 tasks

alamb mentioned this pull request May 13, 2024

DataFusion weekly project plan (Andrew Lamb) - May 13, 2024 #10482

Closed

8 tasks

github-actions bot added the Stale PR has not had any activity for some time label Jul 7, 2024

alamb closed this Jul 7, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

RFC: Make it easier to call window functions via expression API (and add example) #6746

RFC: Make it easier to call window functions via expression API (and add example) #6746

alamb commented Jun 22, 2023 •

edited

Loading

alamb Jun 22, 2023

alamb Jun 22, 2023

timsaucer commented May 2, 2024

alamb commented May 3, 2024

timsaucer commented May 3, 2024

alamb commented May 6, 2024

alamb commented May 6, 2024

timsaucer commented May 7, 2024

alamb commented May 7, 2024

github-actions bot commented Jul 7, 2024

alamb commented Jul 7, 2024

timsaucer commented Jul 8, 2024

alamb commented Jul 9, 2024

RFC: Make it easier to call window functions via expression API (and add example) #6746

RFC: Make it easier to call window functions via expression API (and add example) #6746

Conversation

alamb commented Jun 22, 2023 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

alamb Jun 22, 2023

Choose a reason for hiding this comment

alamb Jun 22, 2023

Choose a reason for hiding this comment

timsaucer commented May 2, 2024

alamb commented May 3, 2024

timsaucer commented May 3, 2024

alamb commented May 6, 2024

alamb commented May 6, 2024

timsaucer commented May 7, 2024

alamb commented May 7, 2024

github-actions bot commented Jul 7, 2024

alamb commented Jul 7, 2024

timsaucer commented Jul 8, 2024

alamb commented Jul 9, 2024

alamb commented Jun 22, 2023 •

edited

Loading